import sys; sys.path.insert(0,'../third_party')
from twisted.spread import pb
from twisted.internet import reactor, defer
from twisted.python import util
import cPickle
from common import crypto
import time

class RPCServer(pb.Root):
    def __init__(self, obj):
        self.obj = obj
        self.port = reactor.listenTCP(obj.rpc_port, pb.PBServerFactory(self))
        reactor.callLater(1, lambda : util.println( "RPC Server listening for requests on port ", self.obj.rpc_port))
        self.context_cache = {}
    def __del__(self):
        self.port.stopListening()
    def remote_handle(self, msg, encrypted=False, pickled_pubkey=None, encrypted_symmkey=None):
        obj = self.obj        
        if encrypted:
            if encrypted_symmkey in self.context_cache:
                encrypt_context,symmkey = self.context_cache[encrypted_symmkey]
            else:
                print 'decrypting symmkey'
                pubkey = cPickle.loads(pickled_pubkey)            
                symmkey = crypto.decrypt_key(encrypted_symmkey, pubkey, obj.crypto_db)
                print 'making encrypt_context'
                encrypt_context = crypto.create_encryption_context(symmkey)
                print 'caching encrypted_symmkey-->symmkey,encrypt_context to save work'
                self.context_cache[encrypted_symmkey] = (encrypt_context,symmkey)
            #print 'decrpyting msg'
            msg = crypto.decrypt(msg, encrypt_context)            
        msg = cPickle.loads(msg)
        method = msg['method']
        args = msg['args']
        del msg['method']        
        if not hasattr(obj, method):
            print '#'*30
            print ("unknown method %s" % method).center(30, '#')
            print '#'*30            
            return
        m = getattr(obj, method)        
        ret = m(*args)
        if isinstance(ret, defer.Deferred):
            #print "the result is a deferred, not serializing, instead sending a remote object"
            return RemoteDeferred(ret)
            
        plain = cPickle.dumps(ret)
        if not encrypted:
            return plain                
        ciphertext = crypto.encrypt(plain, encrypt_context)        
        mac = crypto.hmac(ciphertext, symmkey)
        return (ciphertext, mac)

class _RPCMethod:
    
    def __init__(self, method, client):
        self.method = method
        self.client = client
    def __call__(self, *args,**kwargs):       
        callback = kwargs.get('callback', None)                    
        msg = {'method': self.method, 'args':args}        
        msg = cPickle.dumps(msg)                        
        d = self.client.factory.getRootObject()
        if(self.client.encrypted):
                #print "encrypting the message"            
                msg = crypto.encrypt(msg, self.client.encrypt_context)                        
                d.addCallback(lambda object: 
                        object.callRemote("handle", msg, True, 
                                            self.client.pickled_pubkey, 
                                            self.client.encrypted_symmkey))
                def decrypt_callback(resp):
                    if isinstance(resp, pb.RemoteReference):
                        #print "not decrypting the remote reference object"
                        return RemoteRefWrapper(resp)
                    #print "decrypting the response"
                    ciphertext, receivedMAC = resp                    
                    computedMAC = crypto.hmac(ciphertext, self.client.symmkey)
                    if(receivedMAC != computedMAC):
                        raise Exception("MAC value received from server is incorrect")
                    return cPickle.loads(crypto.decrypt(ciphertext, self.client.encrypt_context))                                            
                d.addCallback(decrypt_callback)
        else:
                print "adding plain handle callback"
                d.addCallback(lambda object:
                                object.callRemote("handle", msg))
                def plain_callback(resp):
                    if isinstance(resp, pb.RemoteReference):
                        print "not unpickling the remote reference object"
                        return RemoteRefWrapper(resp)
                    return cPickle.loads(resp)
                d.addCallback(plain_callback)                 
        if callback:                                   
            d.addCallback(callback)
        else:            
            return d
class RPCClient:
    '''\
    Turns method calls into network messages and returns the output 
    returned by the remote object. See the documentation for RPCServer 
    for an example.
    '''
    encrypted = False
    def __init__(self, host, port, ping_interval=-1):
        factory = RPCClientFactory(ping_interval)       
        self.connector = reactor.connectTCP(host, port, factory)
        self.factory = factory
        self.host = host
        self.port = port
    
    def enable_encryption(self, pubkey, pickled_pubkey=None):
        print 'enable_encryption'
        self.pubkey = pubkey
        self.pickled_pubkey = pickled_pubkey or cPickle.dumps(pubkey)
        self.encrypted = True
        self.symmkey = crypto.random_symmetric_key()
        self.encrypted_symmkey = crypto.encrypt_key(self.symmkey, self.pubkey)
        self.encrypt_context = crypto.create_encryption_context(self.symmkey)
    def __getattr__(self, method):        
        if method in ['mro', 'has_key'] or method.startswith('_'):
            raise AttributeError()
        if method == 'status':
            return self.factory.connected        
        return _RPCMethod(method, self)
    def __repr__(self):        
        return '<RPCClient instance at %s>' % hex(id(self))
    def __del__(self):
        print "disconnecting the rpc client"        
        self.connector.disconnect()        

class RPCClientFactory(pb.PBClientFactory):    
    connected = False
    _setconnect = None
    def __init__(self,ping_interval):
        pb.PBClientFactory.__init__(self,unsafeTracebacks=True)
        self.ping_interval = ping_interval             
    def startedConnecting(self, connector):
        """Called when a connection has been started.

        You can call connector.stopConnecting() to stop the connection attempt.

        @param connector: a Connector object.
        """        
        def nit():
            self.connected = True          
        if not self._setconnect:
            self._setconnect = reactor.callLater(2, nit)
    
    def _disconnected(self, connector, reason):
        self.connected = False        
        if self._setconnect:
            if self._setconnect.active():
                self._setconnect.cancel()
            self._setconnect = None            
        if self.ping_interval != -1:
            reactor.callLater(self.ping_interval, connector.connect)
            pb.PBClientFactory.clientConnectionLost(self, connector, reason, 1)
        else:
            pb.PBClientFactory.clientConnectionLost(self, connector, reason, 0)
    
    def clientConnectionFailed(self, connector, reason):
        """Called when a connection has failed to connect.

        It may be useful to call connector.connect() - this will reconnect.

        @type reason: L{twisted.python.failure.Failure}
        """
        self._disconnected(connector, reason)
    def clientConnectionLost(self, connector, reason):
        """Called when an established connection is lost.

        It may be useful to call connector.connect() - this will reconnect.

        @type reason: L{twisted.python.failure.Failure}
        """
        self._disconnected(connector, reason)
    
class RemoteDeferred(pb.Referenceable):
    def __init__(self, d):        
        self.d = d
    def remote_addCallback(self, obj, *args, **kwargs):
        def _callback(data, *args, **kwargs):
            obj.callRemote("handle", data, *args, **kwargs)
        self.d.addCallback(_callback)
    def remote_addErrback(self, obj, *args, **kwargs):
        def _callback(data, *args, **kwargs):
            obj.callRemote("handle", data, *args, **kwargs)
        self.d.addErrback(_callback)
class LocalFunc(pb.Referenceable):
    def __init__(self, func):        
        self.func = func
    def remote_handle(self, *args, **kwargs):
        self.func(*args, **kwargs)
class RemoteRefWrapper:
    def __init__(self, ref):
        if not isinstance(ref, pb.RemoteReference):
            raise Exception("remoteRef must be a pb.RemoteReference instance")
        self.ref = ref
    def __getattr__(self, method):
        def _method(func, *args, **kwargs):
            self.ref.callRemote(method, LocalFunc(func), *args, **kwargs)
        return _method
        
if __name__=='__main__':
    pass